#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Chart type identification component
"""

import json
from typing import Dict, Any, Optional
from data_engine.models.data_models import Dataset, ChartAnalysis, QueryContext
from data_engine.config import ConfigManager, ChartType
from data_engine.utils.logger import engine_logger
try:
    from config import settings
except ImportError:
    # 如果无法导入，创建一个默认的settings对象
    class DefaultSettings:
        DEBUG = False  # 默认非调试模式
    settings = DefaultSettings()


class ChartIdentifier:
    """图表类型识别器"""
    
    def __init__(self, llm_client):
        """
        初始化图表识别器
        
        Args:
            llm_client: LiteLLM客户端实例
        """
        self.llm_client = llm_client
        self.config_manager = ConfigManager()
    
    async def analyze_chart_requirement(self, dataset: Dataset, 
                                      context: QueryContext) -> ChartAnalysis:
        """
        分析图表需求
        
        Args:
            dataset: 数据集
            context: 查询上下文
            
        Returns:
            ChartAnalysis: 图表分析结果
        """
        if settings.DEBUG:
            engine_logger.debug("\n===== 开始图表需求分析 =====")
            engine_logger.debug(f"分析问题: {context.question}")
        
        engine_logger.info("开始分析图表需求")
        
        try:
            # 构建分析提示词
            prompt = self._build_analysis_prompt(dataset, context)
            
            if settings.DEBUG:
                engine_logger.debug("\n===== 系统提示词 =====")
                engine_logger.debug(prompt)
                engine_logger.debug("\n===== 调用LLM进行分析 =====")
            
            # 调用LLM进行分析
            response = await self.llm_client.acompletion(
                messages=[{"role": "user", "content": prompt}]
            )
            
            # 解析并验证结果
            result_text = self.llm_client.get_response_content(response).strip()
            
            if settings.DEBUG:
                engine_logger.debug(f"LLM原始分析结果:\n{result_text}")
            
            chart_analysis = self._parse_and_validate_result(result_text, dataset)
            
            if settings.DEBUG:
                engine_logger.debug(f"解析后的分析结果: chart_type={chart_analysis.chart_type}, x_field={chart_analysis.x_field}, y_field={chart_analysis.y_field}, group_by={chart_analysis.group_by}, sort_by={chart_analysis.sort_by}")
                if chart_analysis.reasoning:
                    engine_logger.debug(f"分析推理: {chart_analysis.reasoning}")
                engine_logger.debug("===== 图表需求分析完成 =====\n")
            
            engine_logger.info(f"图表分析完成: {chart_analysis.chart_type}, 置信度: {chart_analysis.confidence}")
            return chart_analysis
            
        except Exception as e:
            engine_logger.error(f"图表分析失败: {str(e)}")
            return self._get_fallback_analysis()
    
    def _build_analysis_prompt(self, dataset: Dataset, context: QueryContext) -> str:
        """构建分析提示词"""
        # 获取提示词模板
        template = self.config_manager.get_prompt_template("chart_analysis")
        
        # 准备数据信息
        sample_data = dataset.sample_data.to_string(index=False)
        chart_types_info = self.config_manager.build_chart_types_info()
        
        # 智能构建字段描述
        field_descriptions = []
        for col in dataset.df.columns:
            dtype = str(dataset.df[col].dtype)
            unique_count = dataset.df[col].nunique()
            
            # 基础字段信息
            field_info = f"- {col}: {dtype}, {unique_count} unique values"
            
            # 如果有字段描述，添加到信息中
            if (dataset.field_descriptions and 
                col in dataset.field_descriptions and 
                dataset.field_descriptions[col]):
                field_info += f" ({dataset.field_descriptions[col]})"
            
            field_descriptions.append(field_info)
        
        # 智能组装描述信息
        description_parts = []
        
        # 数据集描述
        if dataset.dataset_description:
            description_parts.append(f"数据集描述: {dataset.dataset_description}")
        
        # 表格描述
        if dataset.table_description:
            description_parts.append(f"表格描述: {dataset.table_description}")
        
        # 如果没有任何描述，添加提示
        if not description_parts:
            description_parts.append("请根据数据字段和样本数据理解数据含义")
        
        # 组装最终描述
        combined_description = "\n".join(description_parts)
        
        return template.format(
            dataset_description=combined_description,
            table_description="",  # 已合并到dataset_description中
            field_descriptions="\n".join(field_descriptions),
            sample_data=sample_data,
            question=context.question
        )
    
    def _parse_and_validate_result(self, result_text: str, dataset: Dataset) -> ChartAnalysis:
        """解析并验证LLM结果"""
        try:
            # 清理JSON格式
            if result_text.startswith('```json'):
                result_text = result_text.replace('```json', '').replace('```', '').strip()
            elif result_text.startswith('```'):
                result_text = result_text.replace('```', '').strip()
            
            data = json.loads(result_text)
            
            if settings.DEBUG:
                engine_logger.debug(f"LLM原始响应: {result_text}")
                engine_logger.debug(f"解析后的数据: {data}")
            
            # 创建分析结果
            chart_analysis = ChartAnalysis(
                chart_type=data.get('chart_type', 'table'),
                need_chart=data.get('need_chart', True),
                x_field=data.get('x_field'),
                y_field=data.get('y_field'),
                y_fields=data.get('y_fields', []),
                name_field=data.get('name_field'),
                value_field=data.get('value_field'),
                group_by=data.get('group_by'),
                filter_condition=data.get('filter_condition'),
                aggregation=data.get('aggregation'),
                sort_by=data.get('sort_by'),
                sort_order=data.get('sort_order', 'asc'),
                limit=data.get('limit'),
                confidence=data.get('confidence', 0.5),
                reasoning=data.get('reasoning', ''),
                series_config=data.get('series_config', {}),
                execution_order=data.get('execution_order', [])
            )
            
            # 验证结果
            return self._validate_analysis_fields(chart_analysis, dataset)
            
        except json.JSONDecodeError:
            engine_logger.warning("JSON解析失败，返回默认分析")
            return self._get_fallback_analysis()
        except Exception as e:
            engine_logger.error(f"解析失败: {str(e)}")
            return self._get_fallback_analysis()
    
    def _validate_analysis_fields(self, analysis: ChartAnalysis, dataset: Dataset) -> ChartAnalysis:
        """验证分析结果字段"""
        # 检查图表类型是否有效
        if not self.config_manager.is_valid_chart_type(analysis.chart_type):
            engine_logger.warning(f"无效的图表类型: {analysis.chart_type}")
            return self._get_fallback_analysis()
        
        # 获取数据列名
        df_columns = dataset.df.columns.tolist()
        
        # 对于表格类型，不需要严格的字段验证
        if analysis.chart_type == 'table':
            return analysis
        
        # 验证字段存在性（但不强制要求）
        self._validate_field_exists(analysis, 'x_field', df_columns)
        self._validate_field_exists(analysis, 'y_field', df_columns)
        self._validate_field_exists(analysis, 'name_field', df_columns)
        self._validate_field_exists(analysis, 'value_field', df_columns)
        
        # 验证y_fields列表
        if analysis.y_fields:
            valid_y_fields = [f for f in analysis.y_fields if f in df_columns]
            analysis.y_fields = valid_y_fields
            engine_logger.info(f"验证后的y_fields: {valid_y_fields}")
        
        # 智能字段检查和修复
        analysis = self._smart_field_validation(analysis, dataset)
        
        return analysis
    
    def _smart_field_validation(self, analysis: ChartAnalysis, dataset: Dataset) -> ChartAnalysis:
        """智能字段验证和修复"""
        df_columns = dataset.df.columns.tolist()
        numeric_cols = dataset.df.select_dtypes(include=['number']).columns.tolist()
        categorical_cols = dataset.df.select_dtypes(include=['object', 'category']).columns.tolist()
        
        chart_type = analysis.chart_type
        
        # 对于折线图和柱状图
        if chart_type in ['line', 'bar', 'area']:
            # 如果有y_fields，优先使用多序列
            if analysis.y_fields and len(analysis.y_fields) > 0:
                # 确保x_field存在
                if not analysis.x_field:
                    # 自动选择合适的x字段
                    if categorical_cols:
                        analysis.x_field = categorical_cols[0]
                    elif df_columns:
                        analysis.x_field = df_columns[0]
                
                engine_logger.info(f"多序列图表: x_field={analysis.x_field}, y_fields={analysis.y_fields}")
                return analysis
            
            # 单序列情况
            if not analysis.x_field and categorical_cols:
                analysis.x_field = categorical_cols[0]
            
            if not analysis.y_field and numeric_cols:
                analysis.y_field = numeric_cols[0]
        
        # 对于饼图
        elif chart_type == 'pie':
            if not analysis.name_field and categorical_cols:
                analysis.name_field = categorical_cols[0]
            
            if not analysis.value_field and numeric_cols:
                analysis.value_field = numeric_cols[0]
        
        # 对于散点图
        elif chart_type == 'scatter':
            if not analysis.x_field and len(numeric_cols) > 0:
                analysis.x_field = numeric_cols[0]
            
            if not analysis.y_field and len(numeric_cols) > 1:
                analysis.y_field = numeric_cols[1]
            elif not analysis.y_field and len(numeric_cols) > 0:
                analysis.y_field = numeric_cols[0]
        
        # 最后检查：如果关键字段仍然缺失，但置信度较高，不降级
        if analysis.confidence > 0.7:
            engine_logger.info(f"高置信度分析({analysis.confidence})，保持图表类型: {chart_type}")
            return analysis
        
        # 低置信度且缺少关键字段时才降级
        required_fields_missing = self._check_critical_fields_missing(analysis)
        if required_fields_missing:
            engine_logger.warning(f"关键字段缺失: {required_fields_missing}")
            return self._create_table_analysis("关键字段缺失")
        
        return analysis
    
    def _check_critical_fields_missing(self, analysis: ChartAnalysis) -> bool:
        """检查是否缺少关键字段"""
        chart_type = analysis.chart_type
        
        if chart_type in ['line', 'bar', 'area']:
            # 多序列优先
            if analysis.y_fields and len(analysis.y_fields) > 0:
                return not analysis.x_field
            # 单序列
            return not (analysis.x_field and analysis.y_field)
        
        elif chart_type == 'pie':
            return not (analysis.name_field and analysis.value_field)
        
        elif chart_type == 'scatter':
            return not (analysis.x_field and analysis.y_field)
        
        return False
    
    def _validate_field_exists(self, analysis: ChartAnalysis, field_name: str, df_columns: list):
        """验证单个字段是否存在"""
        field_value = getattr(analysis, field_name, None)
        if field_value and field_value not in df_columns:
            setattr(analysis, field_name, None)
            engine_logger.warning(f"{field_name} '{field_value}' 不存在")
    
    def _create_table_analysis(self, reason: str) -> ChartAnalysis:
        """创建表格分析结果"""
        return ChartAnalysis(
            chart_type=ChartType.TABLE.value,
            need_chart=False,
            confidence=0.3,
            reasoning=reason
        )
    
    def _get_fallback_analysis(self) -> ChartAnalysis:
        """获取后备分析结果"""
        return ChartAnalysis(
            chart_type=ChartType.TABLE.value,
            need_chart=False,
            confidence=0.1,
            reasoning="分析失败，使用默认表格显示"
        )
    
    def get_chart_recommendation(self, dataset: Dataset) -> Dict[str, Any]:
        """
        获取图表推荐信息
        
        Args:
            dataset: 数据集
            
        Returns:
            Dict[str, Any]: 推荐信息
        """
        numeric_fields = dataset.get_numeric_fields()
        categorical_fields = dataset.get_categorical_fields()
        datetime_fields = dataset.get_datetime_fields()
        
        recommendations = []
        
        # 基于字段类型推荐图表
        if len(numeric_fields) >= 2:
            recommendations.append({
                "chart_type": ChartType.SCATTER.value,
                "reason": f"发现{len(numeric_fields)}个数值字段，适合散点图分析相关性",
                "fields": {"x_field": numeric_fields[0], "y_field": numeric_fields[1]}
            })
        
        if len(categorical_fields) >= 1 and len(numeric_fields) >= 1:
            recommendations.append({
                "chart_type": ChartType.BAR.value,
                "reason": "适合比较不同类别的数值",
                "fields": {"x_field": categorical_fields[0], "y_field": numeric_fields[0]}
            })
            
            recommendations.append({
                "chart_type": ChartType.PIE.value,
                "reason": "适合显示各类别的占比",
                "fields": {"name_field": categorical_fields[0], "value_field": numeric_fields[0]}
            })
        
        if len(datetime_fields) >= 1 and len(numeric_fields) >= 1:
            recommendations.append({
                "chart_type": ChartType.LINE.value,
                "reason": "适合显示时间序列趋势",
                "fields": {"x_field": datetime_fields[0], "y_field": numeric_fields[0]}
            })
        
        return {
            "recommendations": recommendations,
            "field_summary": {
                "numeric_fields": numeric_fields,
                "categorical_fields": categorical_fields,
                "datetime_fields": datetime_fields
            }
        }